{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 164,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import xml.etree.ElementTree as ET\n",
    "import pickle\n",
    "import os\n",
    "from os import listdir, getcwd\n",
    "from os.path import join"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "img_dir = \"gen/jyz\"\n",
    "xml_dir = img_dir\n",
    "txt_dir = img_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 从obj.names中读取名称"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['JYZ']\n"
     ]
    }
   ],
   "source": [
    "classes = []\n",
    "names_file = open(img_dir+'/obj.names')\n",
    "for name in names_file:\n",
    "    classes.append(name.strip( '\\n' ))\n",
    "print(classes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 生成TXT文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def convert(size, box):\n",
    "    dw = 1./(size[0])\n",
    "    dh = 1./(size[1])\n",
    "    x = (box[0] + box[1])/2.0 - 1\n",
    "    y = (box[2] + box[3])/2.0 - 1\n",
    "    w = box[1] - box[0]\n",
    "    h = box[3] - box[2]\n",
    "    x = x*dw\n",
    "    w = w*dw\n",
    "    y = y*dh\n",
    "    h = h*dh\n",
    "    return (x,y,w,h)\n",
    "\n",
    "def convert_annotation( image_id):\n",
    "    in_file = open('%s/%s.xml'%(xml_dir,image_id))\n",
    "    out_file = open('%s/%s.txt'%(txt_dir,image_id), 'w')\n",
    "    tree=ET.parse(in_file)\n",
    "    root = tree.getroot()\n",
    "    size = root.find('size')\n",
    "    w = int(size.find('width').text)\n",
    "    h = int(size.find('height').text)\n",
    "\n",
    "    for obj in root.iter('object'):\n",
    "        difficult = obj.find('difficult').text\n",
    "        cls = obj.find('name').text\n",
    "        if cls not in classes or int(difficult)==1:\n",
    "            continue\n",
    "        cls_id = classes.index(cls)\n",
    "        xmlbox = obj.find('bndbox')\n",
    "        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))\n",
    "        bb = convert((w,h), b)\n",
    "        out_file.write(str(cls_id) + \" \" + \" \".join([str(a) for a in bb]) + '\\n')\n",
    "        \n",
    "    out_file.close()\n",
    "\n",
    "\n",
    "\n",
    "if not os.path.exists(txt_dir):\n",
    "    os.makedirs(txt_dir)\n",
    "for pathAndFilename in glob.iglob(os.path.join(img_dir, \"*.xml\")):\n",
    "    file_name = os.path.basename(pathAndFilename)\n",
    "    convert_annotation(os.path.splitext(file_name)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 生成yolo的其他文件"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 生成train.txt、test.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "生成train.txt、test.txt\n"
     ]
    }
   ],
   "source": [
    "import glob, os\n",
    "\n",
    "\n",
    "# Directory where the data will reside, relative to 'darknet.exe'\n",
    "path_data = img_dir+\"/\"\n",
    "\n",
    "# Percentage of images to be used for the test set\n",
    "percentage_test = 10;\n",
    "\n",
    "# Create and/or truncate train.txt and test.txt\n",
    "file_train = open(path_data+'train.txt', 'w')\n",
    "file_test = open(path_data+'test.txt', 'w')\n",
    "\n",
    "# Populate train.txt and test.txt\n",
    "counter = 1\n",
    "index_test = round(100 / percentage_test)\n",
    "\n",
    "for pathAndFilename in glob.iglob(os.path.join(img_dir, \"*.jpg\")):\n",
    "    title, ext = os.path.splitext(os.path.basename(pathAndFilename))\n",
    "    \n",
    "    if counter == index_test:\n",
    "        counter = 1\n",
    "        file_test.write(path_data + title + '.jpg' + \"\\n\")\n",
    "    else:\n",
    "        file_train.write(path_data + title + '.jpg' + \"\\n\")\n",
    "        counter = counter + 1\n",
    "        \n",
    "file_train.close()\n",
    "file_test.close()\n",
    "print(\"生成train.txt、test.txt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 生成obj.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data_file = open(path_data+'obj.data', 'w')\n",
    "data_file.write(\"classes= \" + str(len(classes)) + \"\\n\")\n",
    "data_file.write(\"train  = \" + path_data + \"train.txt\\n\")\n",
    "data_file.write(\"valid  = \" + path_data+ \"test.txt\\n\")\n",
    "data_file.write(\"names = \" + path_data + \"obj.names\\n\")\n",
    "data_file.write(\"backup = \" + path_data + \"backup/\\n\")\n",
    "\n",
    "data_file.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 生成yolo-obj.cfg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "content = r'''\n",
    "[net]\n",
    "batch=64\n",
    "subdivisions=8\n",
    "height=416\n",
    "width=416\n",
    "channels=3\n",
    "momentum=0.9\n",
    "decay=0.0005\n",
    "angle=0\n",
    "saturation = 1.5\n",
    "exposure = 1.5\n",
    "hue=.1\n",
    "\n",
    "learning_rate=0.0001\n",
    "max_batches = 45000\n",
    "policy=steps\n",
    "steps=100,25000,35000\n",
    "scales=10,.1,.1\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=32\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[maxpool]\n",
    "size=2\n",
    "stride=2\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=64\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[maxpool]\n",
    "size=2\n",
    "stride=2\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=128\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=64\n",
    "size=1\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=128\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[maxpool]\n",
    "size=2\n",
    "stride=2\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=256\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=128\n",
    "size=1\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=256\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[maxpool]\n",
    "size=2\n",
    "stride=2\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=512\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=256\n",
    "size=1\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=512\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=256\n",
    "size=1\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=512\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[maxpool]\n",
    "size=2\n",
    "stride=2\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=1024\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=512\n",
    "size=1\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=1024\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=512\n",
    "size=1\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "filters=1024\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "activation=leaky\n",
    "\n",
    "\n",
    "#######\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "filters=1024\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "filters=1024\n",
    "activation=leaky\n",
    "\n",
    "[route]\n",
    "layers=-9\n",
    "\n",
    "[reorg]\n",
    "stride=2\n",
    "\n",
    "[route]\n",
    "layers=-1,-3\n",
    "\n",
    "[convolutional]\n",
    "batch_normalize=1\n",
    "size=3\n",
    "stride=1\n",
    "pad=1\n",
    "filters=1024\n",
    "activation=leaky\n",
    "\n",
    "[convolutional]\n",
    "size=1\n",
    "stride=1\n",
    "pad=1\n",
    "'''\n",
    "content = content + \"filters=\"+str((5+len(classes))*5)\n",
    "\n",
    "content = content + r'''\n",
    "activation=linear\n",
    "\n",
    "[region]\n",
    "anchors = 1.08,1.19,  3.42,4.41,  6.63,11.38,  9.42,5.11,  16.62,10.52\n",
    "bias_match=1\n",
    "'''\n",
    "content = content + 'classes='+str(len(classes))\n",
    "\n",
    "content = content + r'''\n",
    "coords=4\n",
    "num=5\n",
    "softmax=1\n",
    "jitter=.2\n",
    "rescore=1\n",
    "\n",
    "object_scale=5\n",
    "noobject_scale=1\n",
    "class_scale=1\n",
    "coord_scale=1\n",
    "\n",
    "absolute=1\n",
    "thresh = .6\n",
    "random=0\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data_file = open(path_data+'yolo-obj.cfg', 'w')\n",
    "data_file.write(content)\n",
    "\n",
    "data_file.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n./darknet detector train gen/jyz/obj.data gen/jyz/yolo-obj.cfg darknet19_448.conv.23 -gpus 0,1\\n./darknet detector test gen/jyz/obj.data gen/jyz/yolo-obj.cfg gen/jyz/backup/yolo-obj.backup -gpus 0,1 gen/jyz/jyz_51.jpg\\n\\n'"
      ]
     },
     "execution_count": 175,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "'''\n",
    "./darknet detector train gen/jyz/obj.data gen/jyz/yolo-obj.cfg darknet19_448.conv.23 -gpus 0,1\n",
    "./darknet detector test gen/jyz/obj.data gen/jyz/yolo-obj.cfg gen/jyz/backup/yolo-obj.backup -gpus 0,1 gen/jyz/jyz_51.jpg\n",
    "\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
